#! /usr/bin/python3
# -*- coding: UTF-8 -*-

'''※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※
File Name: train.py
Author: GID5564
Description: 训练模型，保存权重
Version: 1.0
Created Time: 22/04/22-21:18:14
※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※'''
  

# %% 导入必要的包 
import tensorflow as tf
import numpy as np
import pathlib

from model import create_model

def trains(root,epochs=10):
    """
    训练
    :param root: 训练数据根目录
    :param epochs: 训练次数，默认为10次。
    """
    
    # 统计文件夹下的所有图片数量
    data_dir = pathlib.Path(root)
    
    # 从文件夹下读取图片，生成数据集
    train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_dir,color_mode="grayscale",
    image_size=(24, 24),batch_size=32)
    

    # 数据集的分类，对应dataset文件夹下有多少图片分类
    class_names = train_ds.class_names

    # 保存数据集分类
    np.save("class_name.npy", class_names)

    # 数据集缓存处理
    AUTOTUNE = tf.data.experimental.AUTOTUNE
    train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)

    # 创建模型
    model = create_model()

    # 训练模型，epochs=10，所有数据集训练10遍
    model.fit(train_ds,epochs=epochs)
    
    # 保存训练后的权重
    model.save_weights('checkpoint/char_checkpoint')

    
    
if __name__ == "__main__":
    # 图片根目录
    root="./dataset"
    trains(root)

    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    